# Copyright 2021-2023 @ Shenzhen Bay Laboratory &
#                       Peking University &
#                       Huawei Technologies Co., Ltd
#
# This code is a part of MindSPONGE:
# MindSpore Simulation Package tOwards Next Generation molecular modelling.
#
# MindSPONGE is open-source software based on the AI-framework:
# MindSpore (https://www.mindspore.cn/)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ============================================================================
"""
Core engine of MindSPONGE (PyTorch Version)
"""

import os
import time
from typing import Union, List
from collections.abc import Iterable
import collections
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch import Tensor
from torch.optim import Optimizer
from torch.utils.data import Dataset, DataLoader

from .simulation import WithEnergyCell, WithForceCell
from .simulation import RunOneStepCell
from .analysis import AnalysisCell
from ..function import any_not_none, get_arguments, get_tensor
from ..potential import PotentialCell, ForceCell
from ..optimizer import Updater, UpdaterMD
from ..system.molecule import Molecule
from ..metrics import MetricCV, get_metrics
from ..callback import Callback, RunContext, CallbackManager

class Sponge():
    r"""Core engine of MindSPONGE for simulation and analysis (PyTorch Version).

        This Cell is the top-level wrapper for the three modules system (`Molecule`), potential (`PotentialCell`)
        and optimizer (`Optimizer`) in MindSPONGE.

        There are three ways to wraps the modules:

        1)  Wraps `system`, `potential` and `optimizer` directly into `Sponge`.

            Example:
            >>> from sponge import Sponge
            >>> md = Sponge(system, potential, optimizer)

            In this way ordinary simulations can be achieved

        2)  Wrap `system` and `potential` with `WithEnergyCell` first, then wrap `WithEnergyCell` and `optimizer`
            with `Sponge`.

            Example:
            >>> from sponge import WithEnergyCell, Sponge
            >>> sys_with_ene = WithEnergyCell(system, potential)
            >>> md = Sponge(sys_with_ene, optimizer=optimizer)

            In this case, the adjustment of the potential can be achieved by adjusting the `WithEnergyCell`,
            for example by setting the `neighbour_list` and the `bias` in `WithEnergyCell`.

        3)  Wrap `system` and `potential` with `WithEnergyCell` first, then wrap `WithEnergyCell` and `optimizer`
            with `RunOneStepCell`, and finally pass the `RunOneStepCell` into 'Sponge'.

            Example:
            >>> from sponge import WithEnergyCell, RunOneStepCell, Sponge
            >>> sys_with_ene = WithEnergyCell(system, potential)
            >>> one_step = RunOneStepCell(sys_with_ene, optimizer=optimizer)
            >>> md = Sponge(one_step)

            In this case, the adjustment of the force can be achieved by adjusting the `RunOneStepCell`, for example
            by adding a `WithForceCell` to the `RunOneStepCell`.

        For simulations:

            Simulation can be performed by executing the member function `Sponge.run`.

            Example:
                >>> from sponge import Sponge
                >>> md = Sponge(system, potential, optimizer)
                >>> md.run(10000)

        For analysis:

            `Sponge` can also analyse the simulation system by `metrics`. The `metrics` should be a dictionary of
            `Metric` or `Colvar`. The value of the `metrics` can be calculated by executing the member function
            `Sponge.analyse`.

            Example:
            >>> from sponge.colvar import Torsion
            >>> from sponge.colvar import Torsion
            >>> phi = Torsion([4, 6, 8, 14])
            >>> psi = Torsion([6, 8, 14, 16])
            >>> md = Sponge(system, potential, optimizer, metrics={'phi': phi, 'psi': psi})
            >>> metrics = md.analyse()
            >>> for k, v in metrics.items():
            >>>     print(k, v)
            phi [[3.1415927]]
            psi [[3.1415927]]

    Args:

        network (Union[Molecule, WithEnergyCell, RunOneStepCell]):
                                    Cell of the simulation system.

        potential (PotentialCell):  Potential energy. Default: None

        optimizer (Optimizer):      Optimizer. Default: None

        metrics (dict):             A Dictionary of metrics for system analysis. The key type of the `dict` should
                                    be `str`, and the value type of the `dict` should be `Metric` or `Colvar`.
                                    Default: None

        analysis (AnalysisCell):    Analysis network. Default: None

    Supported Platforms:

        ``CPU`` ``GPU``

    """

    def __init__(self,
                 network: Union[Molecule, WithEnergyCell, RunOneStepCell],
                 potential: PotentialCell = None,
                 optimizer: Optimizer = None,
                 metrics: dict = None,
                 update_metrics_pace: int = 1,
                 analysis: AnalysisCell = None,
                 **kwargs
                 ):
        self._kwargs = get_arguments(locals(), kwargs)

        # Initialize parallel settings
        self._parallel_mode = None
        self._device_number = torch.cuda.device_count() if torch.cuda.is_available() else 1
        self._global_rank = 0
        self._parameter_broadcast = False
        self._create_time = int(time.time() * 1e9)
        self.update_metrics_pace = update_metrics_pace

        self._force_function = None
        self._bias_function = None
        self._potential_function = None
        
        if optimizer is None:
            if potential is not None:
                raise ValueError('When optimizer is None, potential must also be None!')
            # network is RunOneStepCell: Sponge = RunOneStepCell
            self._simulation_network: RunOneStepCell = network
            self._system_with_energy: WithEnergyCell = self._simulation_network.system_with_energy
            self._system_with_force: WithForceCell = self._simulation_network.system_with_force
            self._optimizer: Optimizer = self._simulation_network.optimizer
            self._system: Molecule = self._simulation_network.system

            if self._system_with_energy is not None:
                self._potential_function: PotentialCell = self._system_with_energy.potential_function
                self._bias_function = self._system_with_energy.bias_function
            
            if self._system_with_force is not None:
                self._force_function: ForceCell = self._system_with_force.force_function
        else:
            self._system_with_force = None
            self._optimizer = optimizer
            if potential is None:
                # network is WithEnergyCell: Sponge = WithEnergyCell + optimizer
                self._system_with_energy: WithEnergyCell = network
                self._simulation_network = RunOneStepCell(
                    energy=self._system_with_energy, optimizer=self._optimizer)
                self._system: Molecule = self._system_with_energy.system
                self._potential_function: PotentialCell = self._system_with_energy.potential_function
                self._bias_function = self._system_with_energy.bias_function
            else:
                # network is system: Sponge = system + potential + optimizer
                self._system: Molecule = network
                self._system_with_energy = WithEnergyCell(self._system, potential)
                self._simulation_network = RunOneStepCell(
                    energy=self._system_with_energy, optimizer=self._optimizer)
                self._potential_function: PotentialCell = potential
                self._bias_function = None
        # Set up metrics
        self._metrics = get_metrics(metrics)
        self._clear_metrics()

        # Set up analysis
        self._analysis = analysis
        if self._analysis is not None:
            self._analysis.set_system(self._system)
        
        if self._optimizer is not None:
            self.use_updater = None
            if isinstance(self._optimizer, Updater):
                self.use_updater = True
            self.time_step = self._optimizer.param_groups[0]['lr']
        else:
            self.use_updater = None
            self.time_step = 0
    
    def to(self, device):
        visited = set()  # 用于记录已经处理过的对象

        def recursive_to(obj):
            if id(obj) in visited:  # 如果对象已经处理过，直接返回
                return
            visited.add(id(obj))  # 记录当前对象

            if obj is None:
                return
            if isinstance(obj, torch.Tensor):
                obj.data = obj.data.to(device)
                return
            if isinstance(obj, dict):
                for v in obj.values():
                    recursive_to(v)
                return
            if isinstance(obj, (list, set)):
                for v in obj:
                    recursive_to(v)
                return
            if hasattr(obj, '__dict__'):
                for v in obj.__dict__.values():
                    recursive_to(v)
                return

        recursive_to(self.__dict__)
        return self
    
    @property
    def energy_names(self) -> List[str]:
        """Get energy term names"""
        if self._system_with_energy is not None:
            return self._system_with_energy.energy_names
        return []
    
    @property
    def bias_names(self) -> List[str]:
        """Get bias term names"""
        if self._system_with_energy is not None:
            if self._system_with_energy.bias_names is not None:
                return self._system_with_energy.bias_names
        return []

    @property
    def num_energies(self) -> int:
        """Get number of energy terms"""
        if self._system_with_energy is not None:
            return self._system_with_energy.num_energies
        return 0

    @property
    def num_biases(self) -> int:
        """Get number of bias terms"""
        if self._system_with_energy is not None:
            return self._system_with_energy.num_biases
        return 0

    def update_neighbour_list(self, steps):
        """update neighbour list"""
        self._simulation_network.update_neighbour_list(steps)
        return self

    def update_bias(self, step: int):
        """update bias potential"""
        self._simulation_network.update_bias(step)

    def update_wrapper(self, step: int):
        """update energy wrapper"""
        self._simulation_network.update_wrapper(step)

    def update_modifier(self, step: int):
        """update force modifier"""
        self._simulation_network.update_modifier(step)

    def change_optimizer(self, optimizer: Optimizer):
        """Change optimizer"""
        if not isinstance(optimizer, Optimizer):
            raise TypeError(f'The optimizer must be an instance of torch.optim.Optimizer '
                          f'but got: {type(optimizer)}')

        self._optimizer = optimizer

        if isinstance(self._optimizer, Updater):
            self.use_updater = True
        else:
            self.use_updater = False

        # 获取学习率 - PyTorch方式
        lr = self._optimizer.param_groups[0]['lr']
        self.time_step = lr

        self._simulation_network = RunOneStepCell(
            energy=self._system_with_energy, optimizer=self._optimizer)
        self._simulation_network.set_pbc_grad(self.use_updater)
        self.to(self._system.coordinate.device)

        return self

    def change_potential(self, potential: PotentialCell):
        """Change potential"""
        if not isinstance(potential, PotentialCell):
            raise TypeError(f'The potential must be an instance of PotentialCell '
                          f'but got: {type(potential)}')

        self._potential_function = potential
        self._system_with_energy = WithEnergyCell(self._system, self._potential_function, self._bias_function)
        self._simulation_network = RunOneStepCell(
            energy=self._system_with_energy, optimizer=self._optimizer)
        self._simulation_network.set_pbc_grad(self.use_updater)

        return self

    def calc_energy(self) -> Tensor:
        """Calculate total energy"""
        if self._system_with_energy is not None:
            return self._system_with_energy()
        return get_tensor(0.0)

    def calc_energies(self) -> Tensor:
        """Calculate energy terms"""
        if self._system_with_energy is not None:
            return self._system_with_energy.energies()
        return get_tensor([])

    def calc_biases(self) -> Tensor:
        """Calculate bias terms"""
        if self._system_with_energy is not None:
            return self._system_with_energy.biases()
        return get_tensor([])

    def run(self,
            steps: int,
            callbacks: Union[Callback, List[Callback]] = None,
            dataset: Dataset = None,
            ):
        """Run simulation"""
        return self._run(steps, callbacks, dataset)

    def calc_potential(self) -> Tensor:
        """calculate and return the potential energy"""
        if self._system_with_energy is None:
            return 0
        return self._system_with_energy()

    def get_energy(self) -> Tensor:
        """Tensor of total potential energy with shape (B, 1). Data type is float."""
        return self._simulation_network.get_energy()

    def get_energies(self) -> Tensor:
        """Tensor of potential energies with shape (B, U). Data type is float."""
        return self._simulation_network.get_energies()

    def get_biases(self) -> Tensor:
        """Tensor bias potential energies with shape (B, V). Data type is float."""
        return self._simulation_network.get_biases()

    def get_bias(self) -> Tensor:
        """Tensor of total bias potential with shape (B, 1). Data type is float."""
        return self._simulation_network.get_bias()

    def analyse(self,
                dataset: Dataset = None,
                callbacks: Union[Callback, List[Callback]] = None,
                ):
        """Run analysis"""
        if self._metrics is None and self._analysis is None:
            raise ValueError('No metrics or analysis network is set!')
        
        # 创建RunContext并初始化
        run_context = RunContext()
        run_context.update_basic_info(
            network=self._simulation_network,
            optimizer=self._optimizer,
            dataset=dataset,
            parallel_mode=self._parallel_mode,
            device_number=self._device_number
        )
        
        # 转换callbacks为列表
        callbacks = self._transform_callbacks(callbacks)
        run_context.list_callback = callbacks
        
        # 创建callback管理器
        callback_manager = CallbackManager(callbacks)
        
        # 开始分析
        callback_manager.begin(run_context)
        
        if dataset is not None:
            dataloader = DataLoader(dataset)
            for data in dataloader:
                callback_manager.step_begin(run_context)
                
                if self._metrics is not None:
                    metrics = {}
                    for name, metric in self._metrics.items():
                        metrics[name] = metric(data)
                    self.update_metrics(metrics)
                    
                if self._analysis is not None:
                    self._analysis(data)
                    
                callback_manager.step_end(run_context)
        else:
            callback_manager.step_begin(run_context)
            
            if self._metrics is not None:
                metrics = {}
                for name, metric in self._metrics.items():
                    metrics[name] = metric(self._system)
                self.update_metrics(metrics)
                
            if self._analysis is not None:
                self._analysis(self._system)
                
            callback_manager.step_end(run_context)
            
        callback_manager.end(run_context)
        
        return self._get_metrics()

    @staticmethod
    def _transform_callbacks(callbacks: Union[Callback, List[Callback], None]) -> List[Callback]:
        """转换callbacks为列表格式"""
        if callbacks is None:
            return []
        if isinstance(callbacks, Callback):
            return [callbacks]
        if isinstance(callbacks, list):
            return callbacks
        raise TypeError(f'Callbacks must be Callback or List[Callback], but got: {type(callbacks)}')

    def _run(self,
             steps: int,
             callbacks: Union[Callback, List[Callback]] = None,
             dataset: Dataset = None,
             ):
        """运行模拟"""
        if steps <= 0:
            raise ValueError(f'Steps must be greater than 0, but got: {steps}')      
        # 创建RunContext并初始化
        run_context = RunContext()
        run_context.update_basic_info(
            network=self._simulation_network,
            optimizer=self._optimizer,
            dataset=dataset,
            parallel_mode=self._parallel_mode,
            device_number=self._device_number,
            bias_names=self.bias_names,
            energy_names=self.energy_names,
            coordinate=self._system.coordinate,
            pbc_box=self._system.pbc_box,
            potential_function=self._potential_function,
            bias_function=self._system_with_energy.bias_function if self._system_with_energy is not None else None,
            force_function=self._force_function,
        )
        run_context.update_progress(
            total_steps=steps,
            cur_step=0,
        )
        
        if self._metrics is not None:
            run_context.update_basic_info(metrics=self._metrics)        
        
        # 转换callbacks为列表
        callbacks = self._transform_callbacks(callbacks)
        run_context.list_callback = callbacks
            
        # 创建callback管理器
        callback_manager = CallbackManager(callbacks)
        
        # 开始运行
        self._run_process(steps, callback_manager, run_context)
        
    def _run_process(self,
                     steps: int,
                     callback_manager: CallbackManager,
                     run_context: RunContext,
                     ):
        """运行过程实现"""
        def _one_step():
            """执行一步模拟"""
            # 更新物理量
            coordinate = self._system.get_coordinate()
            pbc_box = self._system.get_pbc_box()

            callback_manager.step_begin(run_context)
                        
            # 运行一步
            self._simulation_network()

            run_context.update_physical_quantities(
                potential=self.get_energy(),
                energies=self.get_energies(),
                bias=self.get_bias(),
                biases=self.get_biases(),
                volume=self._system.get_volume(),
                coordinate=coordinate,
                pbc_box=pbc_box,
                image=self._system.image,
            )

            if self.use_updater:
                run_context.update_physical_quantities(
                    velocity=self._optimizer.velocity,
                    kinetics=torch.sum(self._optimizer.kinetics, dim=-1),
                    temperature=self._optimizer.temperature,
                    pressure=self._optimizer.pressure.mean(dim=-1) if self._optimizer.pressure is not None else self._optimizer.pressure,
                )
            if self._metrics is not None:
                self.update_metrics(run_context.cur_step)
                run_context.update_metrics(metrics=self._metrics)
                
            # 更新偏置和包装器
            self.update_bias(run_context.cur_step)
            self.update_wrapper(run_context.cur_step)
            self.update_modifier(run_context.cur_step)
            self.update_neighbour_list(run_context.cur_step)
            
            # 更新进度
            run_context.cur_step += 1
            run_context.cur_time += self.time_step
            
            callback_manager.step_end(run_context)
            
            return run_context.get_stop_requested()
            
        # 开始运行
        callback_manager.begin(run_context)
        
        for step in range(run_context.total_steps):
            run_context.cur_step = step
            callback_manager.step_begin(run_context)
            
            if _one_step():
                break
                
        callback_manager.end(run_context)
        
    def _clear_metrics(self):
        """Clear metrics"""
        if self._metrics is not None:
            for metric in self._metrics.values():
                metric.clear()

    def update_metrics(self,step):
        """Update metrics"""
        if self._metrics is not None and step % self.update_metrics_pace == 0:
            for name, metric in self._metrics.items():
                metric.update(self._system.coordinate,self._system.pbc_box)
        
    def _get_metrics(self):
        """Get metrics"""
        if self._metrics is not None:
            metrics = {}
            for name, metric in self._metrics.items():
                metrics[name] = metric.eval()
            return metrics
        return None

    @property
    def create_time(self):
        """Get create time"""
        return self._create_time
